#!/usr/bin/env python

# 通过 tblastn 的结果 获得当前阅读框下的gff 内容

import argparse
import sys


parser = argparse.ArgumentParser(
    description='''从tblastn 结果 提取 augustus 的gff3 对注释进行补充
    用法:
    fetch_gff_from_tblastn_result.py -i augustus.gff3 -b homo_tblastn -o homo_target.gff3''')

parser.add_argument('-i',
                help='gff3文件')

parser.add_argument('-b',
                help='tblastn结果')

parser.add_argument('-o',
                help='输出文件')


args = parser.parse_args()

if not args.i or not args.b or not args.o:
    parser.print_help()
    sys.exit()


infile = args.b

gff3_file = args.i

outfile = args.o




dic = {}

with open(infile) as fila:
    for i in fila:
        k = i.strip().split('\t')
        if len(k)>2 and k[0] not in dic:
            dic[k[0]] = k

obtained_lista = []
#print(dic)
# 读取gff3 文件

inhandle = open(gff3_file, 'r')


outfile = open(outfile,'w')

t = 0

with open(gff3_file) as fila:
    for i in fila:
        k = i.strip().split('#')[0].split('\t')
        if len(k)>5:

            #print(k)
            if k[2] == 'gene':
                t = 0
                strand = k[6]
                scaf = k[0]
                start = int(k[3])
                end = int(k[4])

                for j in dic:
                    target_scaf = dic[j][1]
                    if int(dic[j][8]) > int(dic[j][9]):
                        target_end = int(dic[j][8])
                        target_start = int(dic[j][9])
                        target_strand = '-'
                    else:
                        target_end = int(dic[j][9])
                        target_start = int(dic[j][8])
                        target_strand = '+'
                    if (scaf == target_scaf and strand == target_strand):

                        if ((start > target_start and start < target_end) or (end < target_end and end > target_start) or (end > target_end and start < target_start)):
                            obtained_lista.append(j)
                            #print(1)
                            outfile.write(i.strip()+'\n')
                            t = 1
                            #print(k,t)
                            break
                        else:
                            t = 0

            else:

                if t == 1:
                    #print(k)
                    outfile.write(i.strip()+'\n')


outfile.close()


if len(obtained_lista)!=len(dic.keys()):
    print('not obtain in:')
    for i in dic:
        if i not in obtained_lista:

            print(i)
